import pandas as pd
from sklearn.model_selection import train_test_split
import pickle

def load_data(seed=2):
    dummies = ['sex', 'smoker', 'region']
    df = pd.read_csv('data/train.csv')
    for i in dummies:
        df = df.join(pd.get_dummies(df[i], prefix=i))
    df = df.drop(dummies, axis=1)
    X = df.drop(['charges'], axis=1)
    y = df[['charges']]
    train_x, valid_x, train_y, valid_y = train_test_split(X, y, test_size=0.1, random_state=seed)
    return train_x, train_y, valid_x, valid_y

def load_test_data():
    x = pd.read_csv('data/test_sample.csv')
    dummies = ['sex', 'smoker', 'region']
    df = x.copy()
    for i in dummies:
        df = df.join(pd.get_dummies(df[i], prefix=i))
    df = df.drop(dummies, axis=1)
    return x, df

def load_model(path):
    print('load model~~~')
    fp = open(path, 'rb')
    model = pickle.load(fp)
    fp.close()
    return model

def save_model(model, path):
    print('save model!!!!!!!!!!!!!!!!!!!!!!!!!!')
    fp = open(path, 'wb')
    pickle.dump(model, fp)
    fp.close()
